package cn.bugstack.chatglm.session.defaults;

import cn.bugstack.chatglm.executor.Executor;
import cn.bugstack.chatglm.model.*;
import cn.bugstack.chatglm.session.Configuration;
import cn.bugstack.chatglm.session.OpenSession;
import okhttp3.sse.EventSource;
import okhttp3.sse.EventSourceListener;

import java.io.IOException;
import java.util.HashMap;

public class DefaultOpenSession implements OpenSession {

    private Configuration configuration;
    private HashMap<Model, Executor> executorGroup;
    public DefaultOpenSession(Configuration configuration,HashMap<Model, Executor> executorGroup) {
        this.configuration = configuration;
        this.executorGroup = executorGroup;
    }
    @Override
    public ChatCompletionResponse chatCompletion(ChatCompletionRequest chatCompletionRequest) throws IOException {
        Executor executor = executorGroup.get(Model.getModel(chatCompletionRequest.getModel()));
        if (executor == null){
            throw new RuntimeException("目前还未支持此大模型");
        }
        return executor.chatCompletion(chatCompletionRequest);
    }

    @Override
    public EventSource completions(ChatCompletionRequest chatCompletionRequest, EventSourceListener eventSourceListener) throws Exception {
        Executor executor = executorGroup.get(chatCompletionRequest.getAiModel());
        if (executor == null){
            throw new RuntimeException("目前还未支持此大模型");
        }
        return executor.completions(chatCompletionRequest,eventSourceListener);
    }

    @Override
    public ImageCompletionResponse genImages(ImageCompletionRequest imageCompletionRequest) throws IOException {
        Executor executor = executorGroup.get(imageCompletionRequest.getAiModel());
        if (executor == null){
            throw new RuntimeException("目前还未支持此大模型");
        }
        return executor.genImages(imageCompletionRequest);
    }

    @Override
    public VideoCompletionResponse genVideos(VideoCompletionRequest videoCompletionRequest) throws Exception {
        Executor executor = executorGroup.get(videoCompletionRequest.getAiModel());
        if (executor == null){
            throw new RuntimeException("目前还未支持此大模型");
        }
        return executor.genVideos(videoCompletionRequest);
    }
}
